#pragma once

#include "bspline_opt.h"
#include <mtx/dynamic.h>
#include <mtx/algorithm.h>
//debug
//#include <mtx/utility.h>
//#include <iostream>

namespace outline {

///abstract class representing a variable or set of variables in bspline
template <typename R>
class bspline_opt_base<R>::variable_base {
public:
	///the dimension of variables' set
	virtual uint size() const = 0;
	///this method may modify goal point (to which given time should be pulled),
	///it is used for vixed control points, which do not have corresponding variables
	virtual void goal_modify(bool left_piece, real t, point & goal) const = 0;
};

//////////////////////////
// x & y separated mode //
//////////////////////////

template <typename R>
class bspline_opt_base<R>::variable : public variable_base {
public:
	///the method builds an equation for the variables' set:
	///point of a bezier piece from the left side (if left_piece) or from the right side,
	///corresponsing to the parameter value (t) should coinside with goal point;
	///the equation is formulated as follows:
	///coefs^T * vars = goal,
	///coefs should point on preallocated space of size() real cells
	virtual void eq(bool left_piece, real t, real * coef) const = 0;
	///modifies right-hand-side of the least-square-error equation according 
	///to the inertia and attractor of the variables' set;
	///rhs points on filled space of size() real cells;
	///double type of rhs is a requirement of mtx library
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const = 0;
};

///typical implementation of variable interface
template <typename R>
class bspline_opt_base<R>::variable_impl : public virtual variable 
{ 
public:
	virtual uint size() const
		{ return 1; }
	virtual void goal_modify(bool left_piece, real t, point & goal) const
		{ /* nothing */ }
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const {
		rhs_x[0] += coef[0] * pp_.x;
		rhs_y[0] += coef[0] * pp_.y;
		coef[0] *= tt_;
	}
public:
	variable_impl() : pp_(0, 0), tt_(1)
		{ }
	void set_traction(point const & attractor, real traction) {
		assert (traction >= 0 && traction < 1);
		//attractor_ = attractor;
		//traction_ = traction;
		tt_ = square(traction / (1 - traction));
		pp_ = attractor * tt_;
		tt_ += 1;
	}
private:
	//point attractor_;
	//real traction_;
	point pp_;
	real tt_;
};	

///position of a bezier node point
template <typename R>
class bspline_opt_base<R>::center : public virtual variable 
{ 
public:
	///given solution of equations, extracts position of the node from it
	///double type of x is a requirement of mtx library;
	///d = 0 for x-coordinate, d = 1 for y-coordinate
	virtual real get_position(uint d, double const * x) const = 0;
};

///right tangent of a bezier node
template <typename R>
class bspline_opt_base<R>::tangent : public virtual variable
{ 
public:
	///given solution of equations, extracts tangent of the node from it
	///double type of x is a requirement of mtx library;
	///d = 0 for x-coordinate, d = 1 for y-coordinate
	virtual real get_tangent(uint d, double const * x) const = 0;
};

///tangents of a bezier node
template <typename R>
class bspline_opt_base<R>::tangents : public virtual variable
{ 
public:
	///given solution of equations, extracts both tangents of the node from it
	///double type of x is a requirement of mtx library;
	///d = 0 for x-coordinate, d = 1 for y-coordinate
	virtual void get_tangents(uint d, double const * x, real & left, real & right) const = 0;
};

///all the data of a bezier node
template <typename R>
class bspline_opt_base<R>::node : public virtual variable
{ 
public:
	///given solution of equations, extracts all the data of the node from it,
	///double type of x is a requirement of mtx library;
	///d = 0 for x-coordinate, d = 1 for y-coordinate
	virtual void get_data(uint d, double const * x, real & center, real & left, real & right) const = 0;
};

template <typename R>
class bspline_opt_base<R>::fixed_center : public center {
public:
	virtual uint size() const
		{ return 0; }
	virtual void goal_modify(bool left_piece, real t, point & goal) const {
		if (!left_piece)
			t = 1 - t;
		goal -= (t * t * (3 - 2 * t)) * position_;
	}
	virtual void eq(bool left_piece, real t, real * coef) const
		{ /* nothing */ }
	virtual real get_position(uint d, double const * x) const 
		{ return position_[d]; }
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const 
		{ /* nothing */ }
public:
	fixed_center() { }
	fixed_center(point const & position) :
		position_(position) { }
	void set_position(point const & position) 
		{ position_ = position; }
private:
	point position_;
};

template <typename R>
class bspline_opt_base<R>::mobile_center : public center, public variable_impl {
public:
	virtual void eq(bool left_piece, real t, real * coef) const {
		if (!left_piece)
			t = 1 - t;
		coef[0] = t * t * (3 - 2 * t);
	}
	virtual real get_position(uint d, double const * x) const {
		return x[0];
	}
};

template <typename R>
class bspline_opt_base<R>::arbitrary_tangent : public tangent, public variable_impl {
public:
	virtual void eq(bool left_piece, real t, real * coef) const {
		if (left_piece)
			coef[0] = 0;
		else
			coef[0] = 3 * t * (1 - t) * (1 - t);
	}
	virtual real get_tangent(uint d, double const * x) const {
		return x[0];
	}
};

template <typename R>
class bspline_opt_base<R>::zero_tangent : public tangent {
public:
	virtual uint size() const
		{ return 0; }
	virtual void goal_modify(bool left_piece, real t, point & goal) const
		{ /* nothing */ }
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const 
		{ /* nothing */ }
	virtual void eq(bool left_piece, real t, real * coef) const
		{ /* nothing */ }
	virtual real get_tangent(uint d, double const * x) const 
		{ return 0;	}
};

template <typename R>
class bspline_opt_base<R>::fixed_tangent : public tangent {
public:
	virtual uint size() const
		{ return 0; }
	virtual void goal_modify(bool left_piece, real t, point & goal) const {
		if (!left_piece)
			goal -= (3 * t * (1 - t) * (1 - t)) * tangent_;
	}
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const 
		{ /* nothing */ }
	virtual void eq(bool left_piece, real t, real * coef) const
		{ /* nothing */ }
	virtual real get_tangent(uint d, double const * x) const 
		{ return tangent_[d];	}
public:
	fixed_tangent() { }
	fixed_tangent(point const & tangent) :
		tangent_(tangent) { }
	void set_tangent(point const & tangent) 
		{ tangent_ = tangent; }
private:
	point tangent_;
};

///left and right tangents of a cusp bezier node
template <typename R>
class bspline_opt_base<R>::independent_tangents : public tangents {
public:
	virtual uint size() const
		{ return left_size_ + right_->size(); }
	virtual void goal_modify(bool left_piece, real t, point & goal) const {
		left_->goal_modify(!left_piece, 1 - t, goal);
		right_->goal_modify(left_piece, t, goal);
	}
	virtual void eq(bool left_piece, real t, real * coef) const {
		left_->eq(!left_piece, 1 - t, coef);
		right_->eq(left_piece, t, coef + left_size_);
	}
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const {
		left_->traction_modify(coef, rhs_x, rhs_y);
		right_->traction_modify(coef + left_size_, rhs_x + left_size_, rhs_y + left_size_);
	}
	virtual void get_tangents(uint d, double const * x, real & left, real & right) const {
		left = left_->get_tangent(d, x);
		right = right_->get_tangent(d, x + left_size_);
	}
public:
	independent_tangents() { }
	independent_tangents(tangent_sptr const & l, tangent_sptr const & r) 
		: left_(l), right_(r)
		{ left_size_ = left_->size(); }
	void set_left(tangent_sptr const & l) 
		{ left_ = l; left_size_ = left_->size(); }
	void set_right(tangent_sptr const & r) 
		{ right_ = r; }
private:
	tangent_sptr left_, right_;
	uint left_size_;
};

template <typename R>
class bspline_opt_base<R>::zero_tangents : public tangents {
public:
	virtual uint size() const
		{ return 0; }
	virtual void goal_modify(bool left_piece, real t, point & goal) const
		{ /* nothing */ }
	virtual void eq(bool left_piece, real t, real * coef) const
		{ /* nothing */ }
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const
		{ /* nothing */ }
	virtual void get_tangents(uint d, double const * x, real & left, real & right) const 
		{ left = right = 0; }
};

///tangents of a smooth node with fixed ratio between left and right tangent lengths
template <typename R>
class bspline_opt_base<R>::fixed_relength_tangents : public tangents, public variable_impl {
public:
	virtual void eq(bool left_piece, real t, real * coef) const {
		if (left_piece)
			coef[0] = l_ * 3 * t * t * (1 - t);
		else
			coef[0] = r_ * 3 * t * (1 - t) * (1 - t);
	}
	virtual void get_tangents(uint d, double const * x, real & left, real & right) const {
		left = l_ * x[0];
		right = r_ * x[0];
	}
public:
	void set_ratio(real l, real r)
		{ l_ = l; r_ = r; assert (l != 0 || r != 0); }
private:
	real l_, r_; ///< relative lengths of tangents
};

template <typename R>
class bspline_opt_base<R>::composite_node : public node {
public:
	virtual uint size() const 
		{ return center_size_ + tangents_->size(); }
	virtual void goal_modify(bool left_piece, real t, point & goal) const {
		center_->goal_modify(left_piece, t, goal);
		tangents_->goal_modify(left_piece, t, goal);
	}
	virtual void eq(bool left_piece, real t, real * coef) const {
		center_->eq(left_piece, t, coef);
		tangents_->eq(left_piece, t, coef + center_size_);
	}
	virtual void traction_modify(double * coef, double * rhs_x, double * rhs_y) const {
		center_->traction_modify(coef, rhs_x, rhs_y);
		tangents_->traction_modify(coef + center_size_, rhs_x + center_size_, rhs_y + center_size_);
	}
	virtual void get_data(uint d, double const * x, real & center, real & left, real & right) const {
		center = center_->get_position(d, x);
		tangents_->get_tangents(d, x + center_size_, left, right);
	}
public:
	void set_center(center_sptr const & c) 
		{ center_ = c; center_size_ = center_->size(); }
	void set_tangents(tangents_sptr const & t) 
		{ tangents_ = t; }
private:
	uint center_size_;
	center_sptr center_;
	tangents_sptr tangents_;
};

// spline

///this class describes overdetermined set of equations for spline fitting
template <typename R>
class bspline_opt_base<R>::spline {
public:
	///necessarily call update_smooth_nodes(...) after create(...)
	void create(bool closed, uint nsize,
		std::vector<junc_type> const & junc, controls<R> const & ctrls,
		point_vector const * attractors = 0, real_vector const * tractions = 0);
	///loads new relative lengths in the solvers for smooth nodes
	void update_smooth_nodes(controls<R> const & ctrls,
		point_vector const * attractors = 0, real_vector const * tractions = 0);
	uint nsize() const
		{ return nodes_.size(); }
	uint vsize() const
		{ return x_.size(); }
public: //formulate and solve overdetermined equations
	void eq(uint piece, real t, point goal);
	void eq(uint piece, real t, point, point goal) { eq(piece, t, goal); }
	void eq_solve(controls<R> & ctrls);
	//eq_solve can return summed discrepancy of the solution:
	//chisq = x^T (A^T A) x - 2(A^T b)*x + |b|^2,
	//but it requires additional variables and computing power;
private:
	void eq_reset();
private:
	typedef mtx::dynamic::vector vector;
	typedef mtx::dynamic::matrix matrix;
	struct fixed_relength_item {
		fixed_relength_tangents * t;
		uint l, c, r; ///< left, center and right point indeces in the spline
		bool symmetrical; ///< the node is symmetrical, thus relenths are known in advance
		fixed_relength_item() { }
		fixed_relength_item(fixed_relength_tangents * tt, uint ll, uint cc, uint rr, bool symm) : t(tt), l(ll), c(cc), r(rr), symmetrical(symm) { }
	};
	static void setup_fixed_relength_tangents(fixed_relength_item const & item,
		controls<R> const & ctrls,
		point_vector const * attractors, real_vector const * tractions);
private:
	bool closed_;
	node_vector nodes_;
	uint_vector var_begin_; ///< var_begin_[i] contains start index of node[i] inside variables vector
	matrix m_;
	vector x_, y_;
	std::vector<fixed_relength_item> fixed_relength_vector_; ///< list of all fixed_relength tangents
};

template <typename R>
void bspline_opt_base<R>::spline::create(bool closed, uint nsize,
	std::vector<junc_type> const & junc, controls<R> const & ctrls,
	point_vector const * attractors, real_vector const * tractions)
{
	closed_ = closed;
	uint s = nsize;
	assert (junc.size() >= s);
	if (attractors)
		assert (closed && attractors->size() >= 3*s && tractions->size() >= 3*s ||
				!closed && attractors->size() >= 3*s-2 && tractions->size() >= 3*s-2);
	nodes_.resize(s);
	var_begin_.resize(s+1);
	fixed_relength_vector_.clear();

	uint vsize = 0;
	var_begin_[0] = 0;
	for (uint i = 0; i < s; ++i) 
	{
		uint l = (i == 0) ? 3*s-1 : 3*i-1; //not valid for open splines and i == 0
		uint c = 3*i;
		uint r = 3*i+1; //not valid for open splines and i == s-1

		composite_node * node = new composite_node();
		nodes_[i] = node_sptr(node);

		if (attractors && (*tractions)[c] >= 1) {
			node->set_center(center_sptr(new fixed_center((*attractors)[c])));
		}
		else {
			mobile_center * center = new mobile_center();
			node->set_center(center_sptr(center));
			if (attractors) {
				center->set_traction((*attractors)[c], (*tractions)[c]);
			}
		}

		switch (junc[i]) {
		case junc_flat:
			node->set_tangents(tangents_sptr(new zero_tangents()));
			break;
		case junc_smooth: 
		case junc_smooth_affine: 
		case junc_symmetrical: 
		{
			if (!closed_ && tractions) { 
				if (i == 0 && (*tractions)[r] >= 1) {
					node->set_tangents(tangents_sptr(new independent_tangents(
						tangent_sptr(new zero_tangent()),
						tangent_sptr(new fixed_tangent((*attractors)[r] - (*attractors)[c]))
					)));
					break;
				}
				if (i+1 == s && (*tractions)[l] >= 1) {
					node->set_tangents(tangents_sptr(new independent_tangents(
						tangent_sptr(new fixed_tangent((*attractors)[l] - (*attractors)[c])),
						tangent_sptr(new zero_tangent())
					)));
					break;
				}
			}
			assert (closed_ || i > 0 && i+1 < s);
			if (tractions && (*tractions)[l] >= 1 && (*tractions)[r] >= 1) {
				node->set_tangents(tangents_sptr(new independent_tangents(
					tangent_sptr(new fixed_tangent((*attractors)[l] - (*attractors)[c])),
					tangent_sptr(new fixed_tangent((*attractors)[r] - (*attractors)[c]))
				)));
				break;
			}
			fixed_relength_tangents * tangents = new fixed_relength_tangents();
			node->set_tangents(tangents_sptr(tangents));
			fixed_relength_item item(tangents, l, c, r, junc[i] == junc_symmetrical);
			if (junc[i] == junc_smooth) //not junc_smooth_affine
				fixed_relength_vector_.push_back(item);
			else 
				setup_fixed_relength_tangents(item, ctrls, attractors, tractions);
			break;
		}
		case junc_cusp: {
			independent_tangents * tangents = new independent_tangents();
			node->set_tangents(tangents_sptr(tangents));
			if (!closed_ && i == 0)
				tangents->set_left(tangent_sptr(new zero_tangent()));
			else {
				if (attractors && (*tractions)[l] >= 1)
					tangents->set_left(tangent_sptr(new fixed_tangent((*attractors)[l] - (*attractors)[c])));
				else {
					arbitrary_tangent * left = new arbitrary_tangent();
					tangents->set_left(tangent_sptr(left));
					if (attractors)
						left->set_traction((*attractors)[l] - (*attractors)[c], (*tractions)[l]);
				}
			}
			if (!closed_ && i == s-1)
				tangents->set_right(tangent_sptr(new zero_tangent()));
			else {
				if (attractors && (*tractions)[r] >= 1)
					tangents->set_right(tangent_sptr(new fixed_tangent((*attractors)[r] - (*attractors)[c])));
				else {
					arbitrary_tangent * right = new arbitrary_tangent();
					tangents->set_right(tangent_sptr(right));
					if (attractors)
						right->set_traction((*attractors)[r] - (*attractors)[c], (*tractions)[r]);
				}
			}
			break;
		}
		default:
			assert (false);
		}

		vsize += nodes_[i]->size();
		var_begin_[i+1] = vsize;
	}

	m_.resize(vsize, vsize);
	x_.resize(vsize);
	y_.resize(vsize);
	eq_reset();
}

template <typename R>
inline void bspline_opt_base<R>::spline::eq_reset()
{
	m_ = 0;
	x_ = 0;
	y_ = 0;
}

template <typename R>
void bspline_opt_base<R>::spline::eq(uint piece, real t, point goal)
{
	assert (piece >= 0);
	assert (closed_ && piece < nsize() || !closed_ && piece+1 < nsize());

	if (piece+1 == nsize()) {
		//last piece in closed spline

		if (piece == 0) {
			//last and the only piece
			uint vs = vsize();
			node & node0 = *nodes_[0];
			assert (node0.size() == vsize());
			real_vector coef0(vs);
			real_vector coef1(vs);

			node0.goal_modify(false, t, goal);
			node0.goal_modify(true, t, goal);

			node0.eq(false, t, &coef0[0]);
			node0.eq(true, t, &coef1[0]);
			for (uint i = 0; i < vs; ++i)
				coef0[i] += coef1[i];

			for (uint i = 0; i < vs; ++i) {
				x_(i) += coef0[i] * goal.x;
				y_(i) += coef0[i] * goal.y;
				for (uint j = i; j < vs; ++j) {
					m_(i, j) += coef0[i] * coef0[j];
				}
			}

			return;
		}

		node & node0 = *nodes_[piece];
		node & node1 = *nodes_[0];
		uint vb = var_begin_[piece];
		uint vm = var_begin_[piece+1];
		uint ve = var_begin_[1];
		assert (node0.size() == vm - vb);
		assert (node1.size() == ve);
		real_vector coef0(vm - vb);
		real_vector coef1(ve);

		node0.goal_modify(false, t, goal);
		node1.goal_modify(true, t, goal);

		node0.eq(false, t, &coef0[0]);
		node1.eq(true, t, &coef1[0]);

		//fill only upper-right part of m_
		for (uint i = 0; i < coef0.size(); ++i) {
			x_(vb + i) += coef0[i] * goal.x;
			y_(vb + i) += coef0[i] * goal.y;
			for (uint j = i; j < coef0.size(); ++j) {
				m_(vb + i, vb + j) += coef0[i] * coef0[j];
			}
		}
		for (uint i = 0; i < coef1.size(); ++i) {
			x_(i) += coef1[i] * goal.x;
			y_(i) += coef1[i] * goal.y;
			for (uint j = i; j < coef1.size(); ++j) {
				m_(i, j) += coef1[i] * coef1[j];
			}
		}
		for (uint i = 0; i < coef1.size(); ++i) {
			for (uint j = 0; j < coef0.size(); ++j) {
				assert (vb + j > i);
				m_(i, vb + j) += coef1[i] * coef0[j];
			}
		}
		return;
	}

	node & node0 = *nodes_[piece];
	node & node1 = *nodes_[piece+1];
	uint vb = var_begin_[piece];
	uint vm = var_begin_[piece+1];
	uint ve = var_begin_[piece+2];
	assert (node0.size() == vm - vb);
	assert (node1.size() == ve - vm);
	real_vector coef(ve - vb);

	node0.goal_modify(false, t, goal);
	node1.goal_modify(true, t, goal);

	node0.eq(false, t, &coef[0]);
	node1.eq(true, t, &coef[vm - vb]);

	//fill only upper-right part of m_
	for (uint i = 0; i < coef.size(); ++i) {
		x_(vb + i) += coef[i] * goal.x;
		y_(vb + i) += coef[i] * goal.y;
		for (uint j = i; j < coef.size(); ++j) {
			m_(vb + i, vb + j) += coef[i] * coef[j];
		}
	}
}

template <typename R>
void bspline_opt_base<R>::spline::eq_solve(controls<R> & ctrls)
{
	assert (closed_ && ctrls.pieces() == nsize() || !closed_ && ctrls.pieces()+1 == nsize());

	uint vsize = this->vsize();
	assert (vsize > 0);
	uint nsize = this->nsize();
	//make m_ symmetric by copying upper-right part into lower-left part
	for (uint i = 0; i < vsize; ++i)
		for (uint j = 0; j < i; ++j)
			m_(i, j) = m_(j, i);

	//modify equations according to the tractions
	std::vector<double> diag(vsize);
	for (uint i = 0; i < vsize; ++i)
		diag[i] = m_(i, i);
	for (uint i = 0; i < nsize; ++i) {
		uint vb = var_begin_[i];
		nodes_[i]->traction_modify(&diag[vb], &x_(0) + vb, &y_(0) + vb);
	}
	for (uint i = 0; i < vsize; ++i)
		m_(i, i) = diag[i];

	//solve equations
	//std::clog << "m:\n" << m_;
	//std::clog << "x:\n" << x_;
	//std::clog << "y:\n" << y_;
	uint_vector pivot(vsize-1);
	LU(m_, pivot);
	sle(m_, x_, pivot);
	sle(m_, y_, pivot);
	//std::clog << "m-lup:\n" << m_;
	//std::clog << "x-sol:\n" << x_;
	//std::clog << "y-sol:\n" << y_;

	//copy result into ctrls
	for (uint i = 0; i < nsize; ++i) 
	{
		uint l = (i == 0) ? 3*nsize-1 : 3*i-1; //not valid for open splines and i == 0
		uint c = 3*i;
		uint r = 3*i+1; //not valid for open splines and i == nsize-1

		uint vb = var_begin_[i];
		point center, left_tan, right_tan;
		nodes_[i]->get_data(0, &x_(0) + vb, center.x, left_tan.x, right_tan.x);
		nodes_[i]->get_data(1, &y_(0) + vb, center.y, left_tan.y, right_tan.y);
		if (closed_ || i > 0)
			ctrls[l] = center + left_tan;
		ctrls[c] = center;
		if (closed_ || i+1 < nsize)
			ctrls[r] = center + right_tan;
	}
	if (closed_)
		ctrls[3*nsize] = ctrls[0];
	eq_reset();
}

template <typename R>
void bspline_opt_base<R>::spline::update_smooth_nodes(controls<R> const & ctrls,
	point_vector const * attractors, real_vector const * tractions)
{
	for (uint i = 0; i < fixed_relength_vector_.size(); ++i) {
		setup_fixed_relength_tangents(fixed_relength_vector_[i],
			ctrls, attractors, tractions);
	}
}

template <typename R>
void bspline_opt_base<R>::spline::setup_fixed_relength_tangents(
	fixed_relength_item const & f,
	controls<R> const & ctrls,
	point_vector const * attractors, real_vector const * tractions)
{
	real jl, jr;
	if (f.symmetrical) {
		jl = -1;
		jr = 1;
	}
	else
		ctrls.calc_tangents_relength(f.l, f.c, f.r, jl, jr, true);
	f.t->set_ratio(jl, jr);

	if (attractors) {
		real t = real(0.5) * ((*tractions)[f.l] + (*tractions)[f.r]);
		if (fabs(jl) > fabs(jr))
			f.t->set_traction(((*attractors)[f.l] - (*attractors)[f.c]) / jl, t);
		else
			f.t->set_traction(((*attractors)[f.r] - (*attractors)[f.c]) / jr, t);
	}
}

} //namespace outline
